import pandas as pd
import numpy as np
from scipy.stats import friedmanchisquare
import argparse
import os

# ----------------------------
# 参数设置
# ----------------------------
parser = argparse.ArgumentParser()
parser.add_argument("--experiment", type=str, default="classification",
                    choices=["classification", "regression"])
parser.add_argument("--sensitive", action="store_true", help="是否使用敏感性分析数据")
args = parser.parse_args()

# ----------------------------
# 模型列表
# ----------------------------
model_cl = ["alex", "res", "mobile", "transformer", "rf", "xgb", "svm"]
model_reg = ["linear", "rf", "gbr", "mlp", "lstm", "rnn", "transformer"]
models = model_cl if args.experiment == "classification" else model_reg

# ----------------------------
# 读取数据
# ----------------------------
df_list = []
for model in models:
    if args.sensitive:
        csv_path = f'results_sensitive/{args.experiment}/{model}/10fold/metrics_{model}.csv'
    else:
        csv_path = f'results/{args.experiment}/{model}/10fold/metrics_{model}.csv'
    df = pd.read_csv(csv_path)
    df = df.drop(columns=["fold","cpa","flops","params"], errors="ignore")
    df_list.append(df)

# 确定指标数量和名称
n_metrics = df_list[0].shape[1]
metric_names = df_list[0].columns.tolist()
n_models = len(models)
n_folds = df_list[0].shape[0]

# ----------------------------
# 计算均值、标准差、95%CI
# ----------------------------
results_dict = {"metric": [], "model": [], "mean": [], "std": [], "ci": [], "p_value": []}

for i, metric in enumerate(metric_names):
    # for idx, df in enumerate(df_list):
    #     print(f"df {idx} shape: {df.shape}")
    #     try:
    #         print(df.iloc[:, i].head())
    #     except Exception as e:
    #         print(f"Error at df {idx}, i={i}: {e}")

    # 取每个模型该指标的10个fold值
    values_per_model = [df.iloc[:, i].values for df in df_list]
    
    # Friedman test
    stat, p = friedmanchisquare(*values_per_model)
    
    # 对每个模型计算 mean/std/CI
    for model_idx, model in enumerate(models):
        vals = values_per_model[model_idx]
        mean = np.mean(vals)
        std = np.std(vals, ddof=1)
        ci = 1.96 * std / np.sqrt(n_folds)
        
        results_dict["metric"].append(metric)
        results_dict["model"].append(model)
        results_dict["mean"].append(mean)
        results_dict["std"].append(std)
        results_dict["ci"].append(ci)
        results_dict["p_value"].append(p)

# ----------------------------
# 保存结果
# ----------------------------
df_results = pd.DataFrame(results_dict)
if args.sensitive:
    save_path = f'analysis/results_sensitive/{args.experiment}/summary_results.csv'
else:
    save_path = f'analysis/results/{args.experiment}/summary_results.csv'
os.makedirs(os.path.dirname(save_path), exist_ok=True)
df_results.to_csv(save_path, index=False)

print(f"结果已保存到 {save_path}")
